import pandas as pd
import redis
import sys  # 新增导入sys模块
from sklearn.ensemble import RandomForestRegressor
import joblib
from sklearn.preprocessing import OneHotEncoder
from sklearn.compose import ColumnTransformer
from sklearn.pipeline import Pipeline
from io import StringIO
from tqdm import tqdm
from sklearn.utils import resample

# 连接Redis获取数据
r = redis.Redis(host='localhost', port=6379, decode_responses=True)
table1_json = r.get('japan_trade_table1')
if not table1_json:  # 新增检查：确保Redis中有数据
    raise ValueError("Redis中没有找到japan_trade_table1数据，请先运行data_processing.py")
    
table1 = pd.read_json(StringIO(table1_json), orient='records')
if table1.empty:  # 新增检查：确保数据不为空
    raise ValueError("从Redis获取的数据为空，请检查数据处理步骤")

table2_json = r.get('japan_trade_table2')
table2 = pd.read_json(StringIO(table2_json), orient='records')

# 模型1：国家级预测
valid_dates_mask = table1['日期'].str.match(r'^\d{4}-\d{2}$')
if not valid_dates_mask.any():  # 新增检查：确保有有效日期
    raise ValueError("没有找到格式正确的日期数据，请检查数据格式")
    
table1 = table1[valid_dates_mask].copy()
table1['年份'] = table1['日期'].str.split('-').str[0].astype(int)
table1['月份'] = table1['日期'].str.split('-').str[1].astype(int)

X1 = table1[['年份', '月份', '国家标签']]
y1_export = table1['出口额']
y1_import = table1['进口额']

preprocessor1 = ColumnTransformer(
    transformers=[
        ('country_encoder', OneHotEncoder(), ['国家标签'])
    ],
    remainder='passthrough'
)

# 在模型训练部分添加进度显示
print("正在训练出口额预测模型...")
model1_export = Pipeline([
    ('preprocessor', preprocessor1),
    ('regressor', RandomForestRegressor(
        n_estimators=100, 
        random_state=42,
        verbose=2  # 修改为更详细的日志级别
    ))
])
with tqdm(total=100, desc='出口额模型训练进度', file=sys.stdout) as pbar:
    model1_export.fit(X1, y1_export)
    pbar.update(100)

print("\n正在训练进口额预测模型...") 
model1_import = Pipeline([
    ('preprocessor', preprocessor1),
    ('regressor', RandomForestRegressor(
        n_estimators=100,
        random_state=42,
        verbose=2
    ))
])
with tqdm(total=100, desc='进口额模型训练进度', file=sys.stdout) as pbar:
    model1_import.fit(X1, y1_import)
    pbar.update(100)

joblib.dump(model1_export, 'model1_export.pkl')
joblib.dump(model1_import, 'model1_import.pkl')

# 模型2：总体预测
valid_dates_mask = table2['日期'].str.match(r'^\d{4}-\d{2}$')
table2 = table2[valid_dates_mask].copy()
table2['年份'] = table2['日期'].str.split('-').str[0].astype(int)
table2['月份'] = table2['日期'].str.split('-').str[1].astype(int)

X2 = table2[['年份', '月份']]
y2_export = table2['出口额']
y2_import = table2['进口额']

model2_export = RandomForestRegressor(n_estimators=100, random_state=42)
model2_export.fit(X2, y2_export)

model2_import = RandomForestRegressor(n_estimators=100, random_state=42)
model2_import.fit(X2, y2_import)

joblib.dump(model2_export, 'model2_export.pkl')
joblib.dump(model2_import, 'model2_import.pkl')

print("模型训练完成，已保存至当前目录")